// Time.cpp : Defines the entry point for the console application.
//

#ifndef __DX11_TS_QUERY__
#define __DX11_TS_QUERY__

#include <d3d11.h>

#include <assert.h>
#include <fstream>
#include <stdarg.h>
#include <string>

std::string getTimeStamp();

#define MAX_QUERY_COUNT 128

class MyTsQuery
{
public:
    MyTsQuery(ID3D11Device* device)
    {
        assert(device);
        m_pDevice = device;
        m_pDevice->AddRef();

        device->GetImmediateContext(&m_pContext);

        D3D11_QUERY_DESC desc;

        desc.Query = D3D11_QUERY_TIMESTAMP_DISJOINT;
        desc.MiscFlags = 0;


        HRESULT hr = m_pDevice->CreateQuery(&desc, &m_tsQueryDisjoint);
        assert(SUCCEEDED(hr));

        desc.Query = D3D11_QUERY_TIMESTAMP;
        for (uint32_t i = 0; i < MaxQueryCount; i++)
        {
            HRESULT hr = m_pDevice->CreateQuery(&desc, &m_tsQuery[i]);
            assert(SUCCEEDED(hr));
        }
    }

    ~MyTsQuery()
    {
        for (uint32_t i = 0; i < MaxQueryCount; i++)
        {
            m_tsQuery[i]->Release();
        }
        m_tsQueryDisjoint->Release();

        m_pContext->Release();
        m_pDevice->Release();
    }

    void Begin()
    {
        assert(stage == 0 || stage == 2);

#ifdef _DEBUG
        stage = 1;
#endif // _DEBUG

        m_pContext->Begin(m_tsQueryDisjoint);
    }

    void QueryTimestamp(const std::string& msg = "")
    {
        assert(stage == 1);
        assert(m_nUseTsCount < MaxQueryCount);

        m_pContext->End(m_tsQuery[m_nUseTsCount]);
        m_tsMsg[m_nUseTsCount] = msg;
        m_nUseTsCount++;
    }

    void End()
    {
        assert(stage == 1);

#ifdef _DEBUG
        stage = 2;
#endif // _DEBUG

        m_pContext->End(m_tsQueryDisjoint);

        HRESULT hr = S_OK;

        D3D11_QUERY_DATA_TIMESTAMP_DISJOINT tsDisjointData;
        do
        {
            hr = m_pContext->GetData(m_tsQueryDisjoint, &tsDisjointData, sizeof(tsDisjointData), 0);
        } while (hr == S_FALSE);
        if (!tsDisjointData.Disjoint)
        {
            assert(m_nUseTsCount > 0);
            UINT64 time0, time1;
            do {
                hr = m_pContext->GetData(m_tsQuery[0], (void*)&time0, sizeof(UINT64), 0);
            } while (hr == S_FALSE);
            assert(hr == S_OK);

            m_nTimeIntervalCount = m_nUseTsCount - 1;
            for (uint32_t i = 0; i < m_nTimeIntervalCount; i++)
            {
                do {
                    hr = m_pContext->GetData(m_tsQuery[i+1], (void*)&time1, sizeof(UINT64), 0);

                } while (hr == S_FALSE);
                assert(hr == S_OK);

                double drawTime = static_cast<double>(time1 - time0) / static_cast<double>(tsDisjointData.Frequency);
                assert(time1 >= time0);

                m_tsTime[i] = drawTime;
                time0 = time1;
            }
        }
        m_nUseTsCount = 0;
    }

    uint32_t GetTimeIntervalCount()
    {
#ifdef _DEBUG
        assert(stage == 2);
#endif // _DEBUG

        return m_nTimeIntervalCount;
    }
    double GetTimeInterval(uint32_t index)
    {
#ifdef _DEBUG
        assert(stage == 2);
#endif // _DEBUG

        return m_tsTime[index]; 
    }

    const std::string& GetTimeIntervalMsg(uint32_t index)
    {
#ifdef _DEBUG
        assert(stage == 2);
#endif // _DEBUG

        return m_tsMsg[index+1];
    }

private:
    uint32_t m_nUseTsCount = 0;
    uint32_t m_nTimeIntervalCount = 0;
    const uint32_t MaxQueryCount = MAX_QUERY_COUNT;
    ID3D11Query* m_tsQueryDisjoint;
    ID3D11Query* m_tsQuery[MAX_QUERY_COUNT];
    double m_tsTime[MAX_QUERY_COUNT] = {};
    std::string m_tsMsg[MAX_QUERY_COUNT];

    ID3D11Device* m_pDevice = nullptr;
    ID3D11DeviceContext* m_pContext = nullptr;

#ifdef _DEBUG
    int stage = 0; // 0 - begin -> 1 - end -> 0
#endif // _DEBUG
};

#endif // !__DX11_TS_QUERY__
